
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.util.HashMap;
import java.util.Map;


public class KnownPlainTextAttack {
//	static final int BlockSize = 8128; 
	//static final int keysize = 52;
//	static int state = 0;//0=Encrypt,1=Decrypt
//	static FileOutputStream output=null;
//	static int test =0;
	//###Attack - AttackCipher methods###//
	/*
	 * a-z: 0-25
	 * A-Z: 26-51
	*/
//	static private int[] aAbBcC = {97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90};
	
	
	public static void main(String [] args)
	{
	/*	
		if (args.length!=6 || args[0].length() <2 || args[1].length() <2 || args[2].length() < 2 || args[3].length()<2 ||args[4].length()<2 ||args[5].length()<2)
		{
			System.out.println("wrong arguments");
			return;
		}
*/
		//args
		String mode = args[0];
		mode = mode.toLowerCase();
		mode = mode.trim();
		
		String messagefile="";
		String key="";
		String IV="";
		String sampleTxt="";
		String sampleChiper="";
		
		if(mode.equals("encryptionb")) 
		{
			state=0;
			messagefile = args[1];
			key = args[2];
			IV= args[3];
			
		}

		else if(mode.equals("decryptionb")){
			state =1;
			sampleTxt = args[1];
			sampleChiper = args[2];
			messagefile = args[3];
			IV= args[4];
		}
		else{
			state=2; // decrypt
			messagefile = args[1];
			key = args[2];
			IV= args[3];
		}
			
		

		//end args
		
		byte [] messageAsBytes = null;
		byte [] IVAsBytes = null;
		byte [] KeyAsBytes=null;
		byte [] sampleTxtAsBytes = null;
		byte [] sampleChiperAsBytes = null;
		Map<Byte,Byte> keyMap = null;
		
		SetupOutputFile(messagefile);
		
		//read args files into an array of bytes
		messageAsBytes = filesnameToByteArray(messagefile);
		IVAsBytes = filesnameToByteArray(IV);
		IVAsBytes = trimIV(IVAsBytes);
		
		if(state==1){
			sampleTxtAsBytes = filesnameToByteArray(sampleTxt);
			sampleChiperAsBytes = filesnameToByteArray(sampleChiper);
		}

		
		if(state!=1){
			KeyAsBytes = filesnameToByteArray(key);
			//end of read args files into an array of bytes
			
			keyMap = keyBytesToMap(KeyAsBytes);
		}

		
		if(state==0)
			createPlainText(messageAsBytes,IVAsBytes,keyMap);
		else if(state ==1)
		{
			AttackCipher(messageAsBytes,IVAsBytes,sampleTxtAsBytes,sampleChiperAsBytes);
		}
		else //state == decrypt
		{
			Map<Byte,Byte> reversedKeyMap=reverseKeyMap(keyMap);
			byte [] decodedText = new byte[messageAsBytes.length];
			decodedText = decodeFileUsingKey(messageAsBytes,IVAsBytes,reversedKeyMap,-1, decodedText);
			PrintBytetofile(decodedText,output);
		}
		
		try 
		{
			output.close();
		} 
		catch (IOException e) 
		{
			System.err.println("Caught IOException: " + e.getMessage());
		}
	}

	
	//open an outputfile according to state (encrypt - cipher, decrypt - txt, attack - key)
	private static void SetupOutputFile(String filename) {
		try {
			int to = filename.indexOf('.');
			if (to==-1)
				to = filename.length();
			
			if (state == 0)
				output = new FileOutputStream(filename.substring(0,to) + ".cipher");
			else if (state == 1)
				output = new FileOutputStream(filename.substring(0,to) + "_key.txt");
			else
				output = new FileOutputStream(filename.substring(0,to) + ".plain.txt");

		} 
		catch (FileNotFoundException e) {
			e.printStackTrace();
		}
	}
	
	//read txt file into an array of bytes 
	private static byte[] filesnameToByteArray(String filename) {
		byte[] ans = null;
		RandomAccessFile f;
		try {
			f = new RandomAccessFile(filename, "r");
			try 
			{
				ans = new byte[(int) f.length()];
			} 
			catch (IOException e) 
			{

				e.printStackTrace();
			}
			try 
			{
				f.read(ans);
			} 
			catch (IOException e) 
			{

				e.printStackTrace();
			}
		} 
		catch (FileNotFoundException e1) 
		{

			e1.printStackTrace();
		}

		return ans;
	}
	
	//IV trim to Block Size
	private static byte[] trimIV(byte[] iVAsBytes) {
		int size = Math.min(iVAsBytes.length, BlockSize);
		byte[] trim = new byte[size];
		for (int i = 0; i < size; i++) {
			trim[i] = iVAsBytes[i];
		}
		return trim;
	}
	
	//create a hash map of keys from the byte array
	private static Map<Byte, Byte> keyBytesToMap(byte[] KeyAsBytes) {

		Map<Byte, Byte> keyMap = new HashMap<Byte, Byte>();

		for (int i = 0; i < keysize; i++) {
			keyMap.put(KeyAsBytes[i * 4], KeyAsBytes[i * 4 + 2]);
		}
		return keyMap;
	}

	//encrypt a txt file
	private static void createPlainText(byte[] messageAsBytes, byte[] iVAsBytes, Map<Byte, Byte> keyMap) {

		int numofiter = messageAsBytes.length / BlockSize;
		if ((double) (messageAsBytes.length) / (double) BlockSize != (double) ((messageAsBytes.length) / BlockSize)) {
			numofiter += 1;
			messageAsBytes = roundMessageArray(messageAsBytes, numofiter);
		}

		byte[] currblock = null;
		for (int i = 0; i < numofiter; i++) {
			currblock = extractCurrBlock(messageAsBytes, i);
			currblock = XOR(currblock, iVAsBytes);
			currblock = encrypt(currblock, keyMap);
			PrintBytetofile(currblock, output);
			iVAsBytes = currblock;

		}
	}
	
	//Retrieve encryption key 
	private static void AttackCipher(byte[] cipher, byte[] IV, byte[] sampleTXT, byte[] sampleCipher) {
		
		Map<Byte, Byte> ans=null;
		//partial key
		Map<Byte, Byte> initialKey = createIntialKey(cipher, IV, sampleTXT, sampleCipher);

		if(initialKey.size()<52)
		{
			 ans = KnownPriorityManager.computeBestKey(initialKey, cipher, IV);
		}
		else
		{
			ans = initialKey;
		}
		PrintBytetofile(KeytoByte(ans), output);
	}


	//Decrypt cipher
	public static byte[] decodeFileUsingKey(byte[] messageAsBytes, byte[] iVAsBytes, Map<Byte, Byte> reversedKeyMap, int currVertexIter, byte[] ans) {

		test++;
		int numofiter = messageAsBytes.length / BlockSize;
		assert (numofiter == messageAsBytes.length / BlockSize);
		if (currVertexIter == -1) { //run on the entire file
			for (int i = 0; i < numofiter; i++) {
				// originalBlock=extractCurrBlock(messageAsBytes, i);
				decrypt(messageAsBytes, reversedKeyMap, i, ans);
				deXOR(messageAsBytes, iVAsBytes, i, ans);

			}
		} 
		else {
			for (int i = (currVertexIter - 1)
					* Letter.numOfBlockInIter; i < numofiter
					&& i < Letter.numOfBlockInIter
							* currVertexIter; i++) {
				// originalBlock=extractCurrBlock(messageAsBytes, i);
				decrypt(messageAsBytes, reversedKeyMap, i, ans);
				deXOR(messageAsBytes, iVAsBytes, i, ans);

			}//end for
		}//end else
		return ans;
	}

	
	//###encrypt - createPlainText methods###//
	//mss doesn't round up with BlockSize
	private static byte[] roundMessageArray(byte[] messageAsBytes,int numofiter) {
		byte [] ans = new byte[numofiter*BlockSize];
		int i;
		assert(ans.length>messageAsBytes.length);
		for(i =0; i<messageAsBytes.length;i++){
			byte tmp = messageAsBytes[i];
			ans[i] = tmp;
		}
		for(;i<ans.length;i++){
			ans[i]=(byte)0;
		}
		return ans;
		
	}
	
	//next block for the iteration
	private static byte[] extractCurrBlock(byte[] messageAsBytes, int i) {
		byte[] ans = new byte[BlockSize];
		for(int inx=0; inx<BlockSize; inx++){
			if(i*BlockSize+inx<messageAsBytes.length)
				ans[inx]=(byte)(messageAsBytes[i*BlockSize+inx]);
			else
			{
				ans[inx] = (byte)(0);
			}
		}
		return ans;
	}
	
	//XOR
	private static byte[] XOR(byte[] messageAsBytes, byte[] iVAsBytes) {
		int size = Math.min(messageAsBytes.length, iVAsBytes.length); // the smallest array
		size = Math.min(size, BlockSize);
		byte[] ans = new byte[size];
		
		for(int inx=0; inx<size; inx++){
			ans[inx]=(byte)(iVAsBytes[inx]^messageAsBytes[inx]);
		
		}
		return ans;
	}
	
	//Encryption
	private static byte[] encrypt(byte[] currblock, Map<Byte, Byte> keyMap) {
		byte[] ans = new byte[BlockSize];
		for(int i=0;i<BlockSize;i++){
			if((currblock[i]>(byte)122 || currblock[i]<(byte)97) && (currblock[i]>(byte)90 || currblock[i]<(byte)65)){
				ans[i]=currblock[i];
			}
			else
			{
			Byte tmp = keyMap.get((Byte)currblock[i]);
				assert(tmp!=null);
					ans[i]=tmp;
			}
		}
		return ans;
		
		
	}
	//###encrypt - createPlainText methods###//
	
	
	//###Decrypt - decodeFileUsingKey methods###//
	//Reverse the key map
	public static Map<Byte, Byte> reverseKeyMap(Map<Byte, Byte> keyMap) {
		
		Map<Byte, Byte> ans = new HashMap<Byte, Byte>();
		for (Map.Entry<Byte, Byte> entry : keyMap.entrySet()) {
			ans.put(entry.getValue(), entry.getKey());
		}
		return ans;
		
	}
	//Decryption
	
	private static void decrypt(byte[] messageAsBytes, Map<Byte, Byte> keyMap,int i, byte[] ans) {
		//byte[] ans = new byte[BlockSize];
		for(int j=0;j<BlockSize;j++){
			byte currbyte = messageAsBytes[i*BlockSize+j];
			if((currbyte>(byte)122 || currbyte<(byte)97) && (currbyte>(byte)90 || currbyte<(byte)65)){
				ans[i*BlockSize+j]=currbyte;
			}
			else{
				Byte tmp = keyMap.get((Byte)currbyte);
				assert(tmp!=null);
				ans[i*BlockSize+j]=tmp;
			}
		}
	}
	
	//XOR
	private static void deXOR(byte[] messageAsBytes, byte[] iVAsBytes,int i, byte[] ans) {
		if(i==0){
			for(int inx=0; inx<BlockSize; inx++){
				ans[inx]=(byte)(iVAsBytes[inx]^ans[inx]);
			
			}
		}
		else
		{
			for(int inx=0; inx<BlockSize; inx++){
				byte val = (byte)(messageAsBytes[(i-1)*BlockSize+ inx]^ans[i*BlockSize+ inx]);
				ans[i*BlockSize+inx]=val;
			
			}
		}

}
	
	//Xor between the sample txt and the iv and learning a partial key
	
	private static Map<Byte,Byte> createIntialKey(byte[] cipher,byte[] IV, byte[] sampleTXT ,byte[] sampleCipher) {

		byte [] xorSample = XOR(sampleTXT,IV);
		Map<Byte,Byte> initialKey = learnKey(xorSample, sampleCipher);
		return initialKey;
	}
	
	//creating the partial key after the xor
	private static Map<Byte, Byte> learnKey(byte[] xorSample, byte[] sampleCipher) {
		
		Map<Byte,Byte> initialKey  = new HashMap<Byte, Byte>();
		byte [] keyArray = new byte[255]; 
		for(int i=0;i<xorSample.length;i++){
			int unsignedByte =  (xorSample[i]) & 0xFF;
			keyArray[unsignedByte] = sampleCipher[i];
		}

		for(int inx = 65; inx < 91; inx++){
			if(keyArray[inx]!=(byte)0){
				initialKey.put((byte)inx, keyArray[inx]);
			}
		}
		
		for(int inx = 97; inx<123; inx++){
			if(keyArray[inx]!=(byte)0){
				initialKey.put((byte)inx, keyArray[inx]);
			}
		}
		
		return initialKey;
	}

	private static void PrintBytetofile(byte[] messageAsBytes, FileOutputStream outfile) {

			try {
				outfile.write(messageAsBytes);
			} catch (IOException e) {
				
				e.printStackTrace();
			}
	}


	public static int [] calcstats(byte[] file, int VertexIter, int [] stats) {

		int start;
		int end;
		if(VertexIter==-1){
			start=0;
			end = file.length;
		}
		else
		{
			start = ((VertexIter-1)*Letter.numOfBlockInIter)*BlockSize ;
			end = VertexIter*Letter.numOfBlockInIter*BlockSize;
		}

		for (int i=start; i<file.length&&i<end;i++){
			int unsignedByte =  (file[i]) & 0xFF;
			if(unsignedByte!=0)
				stats [unsignedByte]++;
			else
			{
				if (i<file.length-BlockSize)
					stats [unsignedByte]++;
			}
		}
		return stats;
	}

	public static byte[] KeytoByte(Map<Byte, Byte> key) {
		byte [] ans = new byte[(52*4)-1];
		int i=0;
		for (int j =0;j<aAbBcC.length;j++){
			ans[i*4+0] = (byte)aAbBcC[j];
			ans[i*4+1] = (byte)(' ');
			ans[i*4+2] = key.get((byte)aAbBcC[j]);
			if(i*4+3<ans.length)
				ans[i*4+3] = (byte)('\n');
			i++;
		}
		return ans;
	}
}
